package com.jsj.subway.ai.largemodel.models;


import com.jsj.subway.ai.common.Role;
import lombok.Data;
import org.apache.commons.lang3.StringUtils;

import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Objects;

@Data
public class AliModelCompletionResponse extends Result.base {


    /**
     * token使用情况
     */
    private Usage usage;

    /**
     * 兼容openAi格式的响应内容
     * 入参result_format=message时候的返回值
     */
    private List<Choice> choices;

    /**
     * 结束原因
     * 正在生成时为null;
     * 生成结束时如果由于停止token导致则为stop;
     * 生成结束时如果因为生成长度过长导致则为length。
     */

    private String id;

    private String model;

    private String object;

    private Long created;

    private Object system_fingerprint;

    /**
     * 本次请求的系统唯一码
     */
    private String request_id;


    /**
     * 兼容openAi格式的响应内容
     */
    @Data
    public static class Choice {
        /**
         * 结束原因，同上
         */
        private String finish_reason;
        /**
         * 响应消息主体
         */
        private Message message;

        private Integer index;

        private Object logprobs;

        private Delta delta;
    }

    /**
     * 本次对话 token使用情况
     *
     * @author Vanffer
     */
    @Data
    public static class Usage {

        private Integer prompt_tokens;

        private Integer completion_tokens;

        private Integer total_tokens;
    }

    @Data
    public static class  Delta {
        private String content;

        private String role;
    }

    public List<Message> toMessage(){
        List<Message> returnMessage = new ArrayList<>();

        List<Choice> choices = this.getChoices();

        for (Choice chatChoice : choices) {
            Message message = chatChoice.getMessage();
            if(Objects.isNull(message)) continue;
            if (!Role.ASSISTANT.getLabel().equals(message.getRole())) continue;

            // 应答完成
            String finishReason = chatChoice.getFinish_reason();
            if (StringUtils.isNoneBlank(finishReason) && "stop".equalsIgnoreCase(finishReason)) {
                return Collections.singletonList(message);
            }
            //可能有多句
            returnMessage.add(message);
        }
        return returnMessage;

    }

    public List<Message> toMessageByStream(){
        List<Message> returnMessage = new ArrayList<>();

        List<Choice> choices = this.getChoices();
        for (Choice chatChoice : choices) {
            Delta delta = chatChoice.getDelta();
            if(Objects.isNull(delta)) continue;
            //if(!Role.ASSISTANT.getLabel().equals(delta.getRole())) continue;
            Message message = new Message();
            message.setRole(delta.getRole());
            if (delta.getContent()!=null){
                message.setContent(delta.getContent());
            }else{
                message.setContent("");
            }


            // 应答完成
            String finishReason = chatChoice.getFinish_reason();
            if (StringUtils.isNoneBlank(finishReason) && "stop".equalsIgnoreCase(finishReason)) {
                return Collections.singletonList(message);
            }

            returnMessage.add(message);
        }
        return returnMessage;
    }
}
